# 实验4的train和test数据划分

import os
from random import shuffle
import random
from util import *  

path_pad = 'E:/data/palm_vein/VERA-PAD/RAW'
path_real = os.path.join(path_pad, 'Real')
path_fake = os.path.join(path_pad, 'Spoof')
            
train_csv_contents = []
test_csv_contents = []
people_ids = list(range(1,50+1))
random.seed(123)  
shuffle(people_ids)
common_train_ids = people_ids[:25] # 编号1~50随机选一半
common_test_ids = people_ids[25:] # 编号1~50随机选另一半

other_ids = list(range(51, 110+1))
shuffle(other_ids)
other_train_ids = other_ids[:30]
other_test_ids = other_ids[30:]

# 读取real图像
real_file_lst = recursive_files(path_real, path_pad)
pid_cid_dict = {}
# real_file_lst_train = []
# real_file_lst_test = []
for file in real_file_lst:
    pid = get_people_id(file)
    if pid in common_train_ids or pid in other_train_ids:
        # real_file_lst_train.append(file)
        train_csv_contents.append([file, 2]) # 2表示真（label会减1）
    else:
        # real_file_lst_test.append(file)
        test_csv_contents.append([file, 2])

# 读取spoof图像
fake_file_lst = recursive_files(path_fake, path_pad)
# fake_file_lst_train = []
# fake_file_lst_test = []
for file in fake_file_lst:
    pid = get_people_id(file)
    if pid in common_train_ids:
        # fake_file_lst_train.append(file)
        train_csv_contents.append([file, 1]) # 1表示假（label会减1）
    else:
        # fake_file_lst_test.append(file)
        test_csv_contents.append([file, 1])

# 保存训练集和测试集的csv文件
write_csv('exp4_vera_pad_train.csv', train_csv_contents)
write_csv('exp4_vera_pad_test.csv', test_csv_contents)
        